#!/usr/bin/env python3

import os
import json
import time
import uuid
import shlex
import psutil
import errno
import argparse
import collections
import subprocess
import logging
import logging.handlers


class H5bench:
    """H5bench benchmark suite."""

    H5BENCH_PATTERNS_WRITE = 'h5bench_write'
    H5BENCH_PATTERNS_READ = 'h5bench_read'
    H5BENCH_EXERCISER = 'h5bench_exerciser'
    H5BENCH_METADATA = 'h5bench_hdf5_iotest'
    H5BENCH_AMREX = 'h5bench_amrex'

    def __init__(self, setup, abort, debug):
        """Initialize the suite."""
        self.LOG_FILENAME = '{}-h5bench.log'.format(setup.replace('.json', ''))

        self.check_parallel()

        self.configure_log(debug)
        self.setup = setup
        self.abort = abort

    def check_parallel(self):
        p = psutil.Process(os.getppid())
        command = p.cmdline()

        # Get user defined shell
        shell = os.environ['SHELL']

        if command[0] != shell:
            print('You should not call MPI directly when running h5bench.')

            # exit(-1)

    def configure_log(self, debug):
        """Configure the logging system."""
        self.logger = logging.getLogger('h5bench')

        if debug:
            self.logger.setLevel(logging.DEBUG)
        else:
            self.logger.setLevel(logging.INFO)

        # Defines the format of the logger
        formatter = logging.Formatter(
            '%(asctime)s %(module)s - %(levelname)s - %(message)s'
        )

        # Configure the log rotation
        handler = logging.handlers.RotatingFileHandler(
            self.LOG_FILENAME,
            maxBytes=268435456,
            backupCount=50,
            encoding='utf8'
        )

        handler.setFormatter(formatter)

        self.logger.addHandler(handler)

        if debug:
            console = logging.StreamHandler()
            console.setFormatter(formatter)

            self.logger.addHandler(console)

    def prepare(self, setup):
        """Create a directory to store all the results of the benchmark."""
        self.directory = setup['directory']

        try:
            # Create a temporary directory to store all configurations
            os.mkdir(self.directory)
        except OSError as exc:
            if exc.errno != errno.EEXIST:
                raise

            self.logger.warning('Base directory already exists: {}'.format(self.directory))

            pass
        except Exception as e:
            self.logger.debug('Unable to create {}: {}'.format(self.directory, e))

        # Check for Lustre support to set the data stripping configuration
        try:
            command = 'lfs getstripe {}'.format(self.directory)

            arguments = shlex.split(command)

            s = subprocess.Popen(arguments, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
            sOutput, sError = s.communicate()

            if s.returncode == 0:
                self.logger.info('Lustre support detected')

                if 'file-system' in setup:
                    if 'lustre' in setup['file-system']:
                        command = 'lfs setstripe'

                        if 'stripe-size' in setup['file-system']['lustre']:
                            command += ' -S {}'.format(setup['file-system']['lustre']['stripe-size'])

                        if 'stripe-count' in setup['file-system']['lustre']:
                            command += ' -c {}'.format(setup['file-system']['lustre']['stripe-count'])

                        command += ' {}'.format(self.directory)

                        self.logger.debug('Lustre stripping configuration: {}'.format(command))

                        arguments = shlex.split(command)

                        s = subprocess.Popen(arguments, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
                        sOutput, sError = s.communicate()

            else:
                self.logger.info('Lustre support not detected')
        except Exception:
            self.logger.info('Lustre support not detected')

    def validate_json(self, setup):
        """Make sure JSON contains all the necessary properties."""
        properties = [
            'mpi',
            'vol',
            'file-system',
            'directory',
            'benchmarks'
        ]

        for p in properties:
            if p not in setup:
                self.logger.critical('JSON configuration file is invalid: "{}" property is missing'.format(p))

                exit(-1)

    def run(self):
        """Run all the benchmarks/kernels."""
        self.logger.info('Starting h5bench Suite')

        try:
            with open(self.setup) as file:
                setup = json.load(file, object_pairs_hook=collections.OrderedDict)
        except Exception:
            self.logger.critical('Unable to find and parse the input configuration file')

            exit(-1)

        self.validate_json(setup)
        self.prepare(setup)

        self.vol_environment = os.environ.copy()
        self.prepare_vol(setup['vol'])

        benchmarks = setup['benchmarks']

        for benchmark in benchmarks:
            name = next(iter(benchmark))
            id = str(uuid.uuid4()).split('-')[0]

            self.logger.info('h5bench [{}] - Starting'.format(name))
            self.logger.info('h5bench [{}] - DIR: {}/{}/'.format(name, setup['directory'], id))

            os.mkdir('{}/{}'.format(setup['directory'], id))

            self.prepare_parallel(setup['mpi'])

            if name == 'write':
                self.run_pattern(id, name, benchmark[name], setup['vol'])
            elif name == 'read':
                self.run_pattern(id, name, benchmark[name], setup['vol'])
            elif name == 'exerciser':
                self.run_exerciser(id, benchmark[name])
            elif name == 'metadata':
                self.run_metadata(id, benchmark[name])
            elif name == 'amrex':
                self.run_amrex(id, benchmark[name], setup['vol'])
            else:
                self.logger('{} - Unsupported benchmark/kernel')

            self.logger.info('h5bench [{}] - Complete'.format(name))

        self.logger.info('Finishing h5bench Suite')

    def prepare_parallel(self, mpi):
        """Prepare the MPI launch command."""
        if 'configuration' in mpi:
            self.mpi = '{} {}'.format(mpi['command'], mpi['configuration'])
        else:
            if mpi['command'] in ['mpirun', 'mpiexec']:
                self.mpi = '{} -np {}'.format(mpi['command'], mpi['ranks'])
            elif mpi['command'] == 'srun':
                self.mpi = '{} --cpu_bind=cores -n {}'.format(mpi['command'], mpi['ranks'])
            else:
                self.logger.warning('Unknown MPI launcher selected!')

                self.mpi = ''

                return

        self.logger.info('Parallel setup: {}'.format(self.mpi))

    def prepare_vol(self, vol):
        """Prepare the environment variables for the VOL."""
        self.vol_environment['ABT_THREAD_STACKSIZE'] = 100000

        if vol is not None:
            if 'library' in vol:
                self.vol_environment['LD_LIBRARY_PATH'] = vol['library']
                self.vol_environment['DYLD_LIBRARY_PATH'] = vol['library']
            if 'path' in vol:
                self.vol_environment['HDF5_PLUGIN_PATH'] = vol['path']

            self.vol_environment['ABT_THREAD_STACKSIZE'] = '100000'

    def enable_vol(self, vol):
        """Enable VOL by setting the connector."""
        if 'connector' in vol:
            self.vol_environment['HDF5_VOL_CONNECTOR'] = vol['connector']

    def disable_vol(self, vol):
        """Disable VOL by setting the connector."""
        if 'HDF5_VOL_CONNECTOR' in self.vol_environment:
            del self.vol_environment['HDF5_VOL_CONNECTOR']

    def reset_vol(self):
        """Reset the environment variables for the VOL."""
        if self.vol_environment is not None:
            if 'LD_LIBRARY_PATH' in self.vol_environment:
                del self.vol_environment['LD_LIBRARY_PATH']
            if 'DYLD_LIBRARY_PATH' in self.vol_environment:
                del self.vol_environment['DYLD_LIBRARY_PATH']
            if 'HDF5_PLUGIN_PATH' in self.vol_environment:
                del self.vol_environment['HDF5_PLUGIN_PATH']
            if 'HDF5_VOL_CONNECTOR' in self.vol_environment:
                del self.vol_environment['HDF5_VOL_CONNECTOR']

            if 'ABT_THREAD_STACKSIZE' in self.vol_environment:
                del self.vol_environment['ABT_THREAD_STACKSIZE']

    def run_pattern(self, id, operation, setup, vol):
        """Run the h5bench_patterns (write/read) benchmarks."""
        self.enable_vol(vol)

        try:
            start = time.time()

            # Define the output file (should be a .h5 file)
            file = '{}/{}'.format(self.directory, setup['file'])
            configuration = setup['configuration']

            configuration_file = '{}/{}/h5bench.cfg'.format(self.directory, id)

            # Create the configuration file for this benchmark
            with open(configuration_file, 'w+') as f:
                for key in configuration:
                    # Make sure the CSV file is generated in the temporary path
                    if key == 'CSV_FILE':
                        configuration[key] = '{}/{}/{}'.format(self.directory, id, configuration[key])

                    f.write('{}={}\n'.format(key, configuration[key]))

            if operation == 'write':
                benchmark_path = self.H5BENCH_PATTERNS_WRITE

            if operation == 'read':
                benchmark_path = self.H5BENCH_PATTERNS_READ

            command = '{} {} {} {}'.format(
                self.mpi,
                benchmark_path,
                configuration_file,
                file
            )

            self.logger.info(command)

            # Make sure the command line is in the correct format
            arguments = shlex.split(command)

            stdout_file_name = '{}/{}/stdout'.format(self.directory, id)
            stderr_file_name = '{}/{}/stderr'.format(self.directory, id)

            with open(stdout_file_name, mode='w') as stdout_file, open(stderr_file_name, mode='w') as stderr_file:
                s = subprocess.Popen(arguments, stdout=stdout_file, stderr=stderr_file, env=self.vol_environment)
                sOutput, sError = s.communicate()

                if s.returncode == 0:
                    self.logger.info('SUCCESS')
                else:
                    self.logger.error('Return: %s (check %s for detailed log)', s.returncode, stderr_file_name)

                    if self.abort:
                        self.logger.critical('h5bench execution aborted upon first error')

                        exit(-1)

            end = time.time()

            self.logger.info('Runtime: {:.7f} seconds (elapsed time, includes allocation wait time)'.format(end - start))
        except Exception as e:
            self.logger.error('Unable to run the benchmark: %s', e)

        self.disable_vol(vol)

    def run_exerciser(self, id, setup):
        """Run the exerciser benchmark."""
        try:
            start = time.time()

            configuration = setup['configuration']

            parameters = []

            # Create the configuration parameter list
            for key in configuration:
                parameters.append('--{} {} '.format(key, configuration[key]))

            command = '{} {} {}'.format(
                self.mpi,
                self.H5BENCH_EXERCISER,
                ' '.join(parameters)
            )

            self.logger.info(command)

            # Make sure the command line is in the correct format
            arguments = shlex.split(command)

            stdout_file_name = '{}/{}/stdout'.format(self.directory, id)
            stderr_file_name = '{}/{}/stderr'.format(self.directory, id)

            with open(stdout_file_name, mode='w') as stdout_file, open(stderr_file_name, mode='w') as stderr_file:
                s = subprocess.Popen(arguments, stdout=stdout_file, stderr=stderr_file, env=self.vol_environment)
                sOutput, sError = s.communicate()

                if s.returncode == 0:
                    self.logger.info('SUCCESS')
                else:
                    self.logger.error('Return: %s (check %s for detailed log)', s.returncode, stderr_file_name)

                    if self.abort:
                        self.logger.critical('h5bench execution aborted upon first error')

                        exit(-1)

            end = time.time()

            self.logger.info('Runtime: {:.7f} seconds (elapsed time, includes allocation wait time)'.format(end - start))
        except Exception as e:
            self.logger.error('Unable to run the benchmark: %s', e)

    def run_metadata(self, id, setup):
        """Run the metadata stress benchmark."""
        try:
            start = time.time()

            # Define the output file (should be a .h5 file)
            file = '{}/{}'.format(self.directory, setup['file'])
            configuration = setup['configuration']

            configuration_file = '{}/{}/hdf5_iotest.ini'.format(self.directory, id)

            # Create the configuration file for this benchmark
            with open(configuration_file, 'w+') as f:
                f.write('[DEFAULT]\n')

                for key in configuration:
                    # Make sure the CSV file is generated in the temporary path
                    if key == 'csv-file':
                        configuration[key] = '{}/{}/{}'.format(self.directory, id, configuration[key])

                    f.write('{} = {}\n'.format(key, configuration[key]))

                f.write('hdf5-file = {}\n'.format(file))

            command = '{} {} {}'.format(
                self.mpi,
                self.H5BENCH_METADATA,
                configuration_file
            )

            self.logger.info(command)

            # Make sure the command line is in the correct format
            arguments = shlex.split(command)

            stdout_file_name = '{}/{}/stdout'.format(self.directory, id)
            stderr_file_name = '{}/{}/stderr'.format(self.directory, id)

            with open(stdout_file_name, mode='w') as stdout_file, open(stderr_file_name, mode='w') as stderr_file:
                s = subprocess.Popen(arguments, stdout=stdout_file, stderr=stderr_file, env=self.vol_environment)
                sOutput, sError = s.communicate()

                if s.returncode == 0:
                    self.logger.info('SUCCESS')
                else:
                    self.logger.error('Return: %s (check %s for detailed log)', s.returncode, stderr_file_name)

                    if self.abort:
                        self.logger.critical('h5bench execution aborted upon first error')

                        exit(-1)

            end = time.time()

            self.logger.info('Runtime: {:.7f} seconds (elapsed time, includes allocation wait time)'.format(end - start))
        except Exception as e:
            self.logger.error('Unable to run the benchmark: %s', e)

    def run_amrex(self, id, setup, vol):
        """Run the AMReX benchmark."""
        self.enable_vol(vol)

        try:
            start = time.time()

            file = '{}/{}'.format(self.directory, setup['file'])
            configuration = setup['configuration']

            configuration_file = '{}/{}/amrex.ini'.format(self.directory, id)

            # Create the configuration file for this benchmark
            with open(configuration_file, 'w+') as f:
                for key in configuration:
                    f.write('{} = {}\n'.format(key, configuration[key]))

                f.write('directory = {}\n'.format(file))

            command = '{} {} {}'.format(
                self.mpi,
                self.H5BENCH_AMREX,
                configuration_file
            )

            self.logger.info(command)

            # Make sure the command line is in the correct format
            arguments = shlex.split(command)

            stdout_file_name = 'stdout'
            stderr_file_name = 'stderr'

            with open(stdout_file_name, mode='w') as stdout_file, open(stderr_file_name, mode='w') as stderr_file:
                s = subprocess.Popen(arguments, stdout=stdout_file, stderr=stderr_file, env=self.vol_environment)
                sOutput, sError = s.communicate()

                if s.returncode == 0:
                    self.logger.info('SUCCESS')
                else:
                    self.logger.error('Return: %s (check %s for detailed log)', s.returncode, stderr_file_name)

                    if self.abort:
                        self.logger.critical('h5bench execution aborted upon first error')

                        exit(-1)

            end = time.time()

            self.logger.info('Runtime: {:.7f} seconds (elapsed time, includes allocation wait time)'.format(end - start))
        except Exception as e:
            self.logger.error('Unable to run the benchmark: %s', e)

        self.disable_vol(vol)


PARSER = argparse.ArgumentParser(
    description='H5bench: a Parallel I/O Benchmark Suite for HDF5: '
)

PARSER.add_argument(
    'setup',
    action='store',
    help='JSON file with the benchmarks to run'
)

PARSER.add_argument(
    '--abort-on-failure',
    action='store_true',
    dest='abort',
    help='Stop h5bench if a benchmark failed'
)

PARSER.add_argument(
    '--debug',
    action='store_true',
    dest='debug',
    help='Enable debug mode'
)

ARGS = PARSER.parse_args()

BENCH = H5bench(ARGS.setup, ARGS.debug, ARGS.abort)
BENCH.run()
